[mlir][IR] Set insertion point when erasing an operation #146955
Conversation
|
@llvm/pr-subscribers-mlir @llvm/pr-subscribers-mlir-core Author: Matthias Springer (matthias-springer) ChangesErasing the operation to which the current insertion point is set, leaves the insertion point in an invalid state. This commit resets the insertion point to the following operation. Note: Requires a traversal of all parent blocks to account for insertion point inside of the erased operation. WIP: 6 test cases have heap-use-after-free errors. There could be something wrong with this PR. Full diff: https://github.com/llvm/llvm-project/pull/146955.diff 1 Files Affected:
diff --git a/mlir/lib/IR/PatternMatch.cpp b/mlir/lib/IR/PatternMatch.cpp
index 1e6084822a99a..95b8d4cac2b29 100644
--- a/mlir/lib/IR/PatternMatch.cpp
+++ b/mlir/lib/IR/PatternMatch.cpp
@@ -152,12 +152,45 @@ void RewriterBase::replaceOp(Operation *op, Operation *newOp) {
eraseOp(op);
}
+/// Returns the given block iterator if it lies within the block `b`.
+/// Otherwise, otherwise finds the ancestor of the given block iterator that
+/// lies within `b`. Returns and "empty" iterator if the latter fails.
+///
+/// Note: This is a variant of Block::findAncestorOpInBlock that operates on
+/// block iterators instead of ops.
+static std::pair<Block *, Block::iterator>
+findAncestorIteratorInBlock(Block *b, Block *itBlock, Block::iterator it) {
+ // Case 1: The iterator lies within the block.
+ if (itBlock == b)
+ return std::make_pair(itBlock, it);
+
+ // Otherwise: Find ancestor iterator. Bail if we run out of parent ops.
+ Operation *parentOp = itBlock->getParentOp();
+ if (!parentOp)
+ return std::make_pair(static_cast<Block *>(nullptr), Block::iterator());
+ Operation *op = b->findAncestorOpInBlock(*parentOp);
+ if (!op)
+ return std::make_pair(static_cast<Block *>(nullptr), Block::iterator());
+ return std::make_pair(op->getBlock(), op->getIterator());
+}
+
/// This method erases an operation that is known to have no uses. The uses of
/// the given operation *must* be known to be dead.
void RewriterBase::eraseOp(Operation *op) {
assert(op->use_empty() && "expected 'op' to have no uses");
auto *rewriteListener = dyn_cast_if_present<Listener>(listener);
+ // If the current insertion point is before/within the erased operation, we
+ // need to adjust the insertion point to be after the operation.
+ if (getInsertionBlock()) {
+ Block *insertionBlock;
+ Block::iterator insertionPoint;
+ std::tie(insertionBlock, insertionPoint) = findAncestorIteratorInBlock(
+ op->getBlock(), getInsertionBlock(), getInsertionPoint());
+ if (insertionBlock && insertionPoint == op->getIterator())
+ setInsertionPointAfter(op);
+ }
+
// Fast path: If no listener is attached, the op can be dropped in one go.
if (!rewriteListener) {
op->erase();
@@ -322,6 +355,11 @@ void RewriterBase::inlineBlockBefore(Block *source, Block *dest,
moveOpBefore(&source->front(), dest, before);
}
+ // If the current insertion point is within the source block, adjust the
+ // insertion point to the destination block.
+ if (getInsertionBlock() == source)
+ setInsertionPoint(dest, getInsertionPoint());
+
// Erase the source block.
assert(source->empty() && "expected 'source' to be empty");
eraseBlock(source);
|
mlir/lib/IR/PatternMatch.cpp
Outdated
| if (getInsertionBlock()) { | ||
| Block *insertionBlock; | ||
| Block::iterator insertionPoint; | ||
| std::tie(insertionBlock, insertionPoint) = findAncestorIteratorInBlock( | ||
| op->getBlock(), getInsertionBlock(), getInsertionPoint()); | ||
| if (insertionBlock && insertionPoint == op->getIterator()) | ||
| setInsertionPointAfter(op); | ||
| } |
There was a problem hiding this comment.
What do you think about simplifying this to just:
if (getInsertionBlock())
if (getInsertionPoint() == op->getIterator())
setInsertionPointAfter(op);
I agree with you that we don't want to make this too expensive. I think the above should cover the cases discussed in #146908 and is a rather very small cost.
I personally do not have expectations of the insertion point still working if the block is being deleted that one is inserting, but I think few users are thinking "My insertion point is before this op and the op is used as an anchor inside the builder to insert new ops, therefore I now need to set the insertion point again after I deleted the op". Even less so in patterns where the framework sets the insertion point for you!
matthias-springer
changed the title
|
Reviewers: Would you rather merge this PR or a bunch of these PRs? |
| if (getInsertionBlock() == op->getBlock() && | ||
| getInsertionPoint() == op->getIterator()) |
There was a problem hiding this comment.
| if (getInsertionBlock() == op->getBlock() && | |
| getInsertionPoint() == op->getIterator()) | |
| if (getInsertionPoint() == op->getIterator()) |
nit: Just comparing iterators ought to be enough
| // If the current insertion point is within the source block, adjust the | ||
| // insertion point to the destination block. | ||
| if (getInsertionBlock() == source) | ||
| setInsertionPoint(dest, getInsertionPoint()); |
There was a problem hiding this comment.
Does this deserve an API documentation update as well?
Erasing the operation to which the current insertion point is set, leaves the insertion point in an invalid state. This commit resets the insertion point to the following operation.
Also adjust the insertion point when inlining a block.